package com.yunchang.spring.visitor.core.handler.response;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.*;
import java.net.URLEncoder;
import java.util.Map;

/**
 * <pre>
 * 处理下载的IResponseHandler。
 * 根据service方法的返回值类型，会进行如下处理：
 *      1.如果返回的是InputStream类型，则会提供下载，但不包含Content-Disposition，无法得知文件名
 *      2.如果返回的File类型，则会提供下载，会包含Content-Disposition，会返回文件名
 *      3.如果返回的Map类型，KEY_INPUT_STREAM必须是文件流InputStream，KEY_FILE_NAME必须是文件名
 *        则会提供下载，会包含Content-Disposition，会返回文件名
 * </pre>
 * Created by jasontujun on 2018/1/26.
 *
 * @see #KEY_INPUT_STREAM
 * @see #KEY_FILE_NAME
 */
public class DownloadHandler implements IResponseHandler {

    /**
     * 指定文件InputStream的key
     */
    public static final String KEY_INPUT_STREAM = "input";

    /**
     * 指定文件名的key
     */
    public static final String KEY_FILE_NAME = "file";

    @Override
    public Object handleResponse(HttpServletRequest request,
                                 HttpServletResponse response,
                                 String service,
                                 String method,
                                 long costTime,
                                 Object result,
                                 Map<String, String> params,
                                 String httpContent) throws Exception {
        InputStream input = null;
        OutputStream output = null;
        String fileName = null;
        if (result instanceof InputStream) {
            input = (InputStream) result;
        } else if (result instanceof File) {
            File file = (File) result;
            input = new FileInputStream(file);
            fileName = file.getName();
        } else if (result instanceof Map) {
            input = (InputStream) ((Map) result).get(KEY_INPUT_STREAM);
            fileName = (String) ((Map) result).get(KEY_FILE_NAME);
        }
        if (input == null) {
            throw new IllegalStateException(String.format("[%s].[%s] not support download!", service, method));
        }
        try {
            output = response.getOutputStream();
            // 设置response的Header
            response.setContentType("application/octet-stream");
            if (fileName != null && fileName.length() != 0) {
                response.addHeader("Content-Disposition", "attachment; filename=" + URLEncoder.encode(fileName, "UTF-8"));
            }
            byte[] data = new byte[1024 * 128];
            int buffer;
            while ((buffer = input.read(data)) > 0) {
                output.write(data, 0, buffer);
            }
            output.flush();
        } finally {
            input.close();
            if (output != null) {
                output.close();
            }
        }
        return null;
    }
}
